/******************************************************************************
 * Gzip-framed ELK Socket Server
 * Protocol (framed mode):
 *   Request/Response Frame:
 *     4 bytes  magic: 'G','Z','1','0'
 *     4 bytes  big-endian unsigned payload length (compressed)
 *     N bytes  gzip-compressed UTF-8 JSON
 *
 * Legacy mode:
 *   If first non-whitespace bytes are not magic, falls back to original
 *   brace-count JSON as implemented in AbstractElkServer#readGraph.
 *
 * Configuration (optional):
 *   -Delk.server.port=5008            (or env ELK_SERVER_PORT)
 *   -Delk.server.host=localhost
 *
 * Loop:
 *   Read frame (or legacy JSON) -> build Elk model -> layout -> send frame (or legacy JSON)
 *
 * Note:
 *   For very large JSON you still build a full string before parsing; move to a streaming
 *   protocol later if needed.
 ******************************************************************************/
package io.typefox.elkserver;

import static java.nio.charset.StandardCharsets.UTF_8;

import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.EOFException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetSocketAddress;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.Channels;
import java.util.concurrent.Executors;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import com.google.gson.Gson;

import org.eclipse.elk.core.RecursiveGraphLayoutEngine;
import org.eclipse.elk.core.util.BasicProgressMonitor;
import org.eclipse.elk.graph.ElkNode;
import org.eclipse.elk.graph.json.ElkGraphJson;

public class GzipFramedSocketElkServer extends AbstractElkServer {

    private static final Logger LOG = Logger.getLogger(GzipFramedSocketElkServer.class.getName());

    private static final byte[] MAGIC = new byte[] { 'G', 'Z', '1', '0' };
    private static final int HEADER_LEN = 8; // 4 magic + 4 length

    private int port;
    private String host;

    public GzipFramedSocketElkServer() {
        this.port = Integer.getInteger("elk.server.port",
                Integer.parseInt(System.getenv().getOrDefault("ELK_SERVER_PORT", "5008")));
        this.host = System.getProperty("elk.server.host",
                System.getenv().getOrDefault("ELK_SERVER_HOST", "localhost"));
    }

    public GzipFramedSocketElkServer(String host, int port) {
        this.host = host;
        this.port = port;
    }

    @Override
    protected void run() throws Exception {
        var executor = Executors.newCachedThreadPool();
        var address = new InetSocketAddress(host, port);
        try (var serverSocket = AsynchronousServerSocketChannel.open().bind(address)) {
            LOG.info("ELK Gzip-Framed Server listening on " + address);
            int clientSeq = 0;
            while (true) {
                AsynchronousSocketChannel channel = serverSocket.accept().get();
                final int clientNo = clientSeq++;
                LOG.info("[" + clientNo + "] Accepted connection");
                executor.submit(() -> {
                    try {
                        handleConnection(channel, clientNo);
                        LOG.info("[" + clientNo + "] Closed connection");
                    } catch (Throwable t) {
                        LOG.log(Level.SEVERE, "[" + clientNo + "] Error: " + t.getMessage(), t);
                    }
                });
            }
        } finally {
            executor.shutdown();
        }
    }

    private void handleConnection(AsynchronousSocketChannel channel, int clientNo) throws Exception {
        InputStream rawIn = Channels.newInputStream(channel);
        OutputStream rawOut = Channels.newOutputStream(channel);
        BufferedInputStream in = new BufferedInputStream(rawIn);
        var out = rawOut;
        var gson = new Gson();
        var engine = new RecursiveGraphLayoutEngine();

        Mode mode = detectMode(in);
        LOG.info("[" + clientNo + "] Mode = " + mode);

        if (mode == Mode.FRAMED) {
            while (true) {
                byte[] jsonBytes = readFramedRequest(in);
                if (jsonBytes == null) {
                    break; // EOF
                }
                LOG.info("[" + clientNo + "] Received framed+gzip JSON (" + jsonBytes.length + " bytes uncompressed)");
                ElkNode graph;
                try {
                    graph = ElkGraphJson.forGraph(new String(jsonBytes, UTF_8)).toElk();
                } catch (Exception parse) {
                    LOG.log(Level.SEVERE, "[" + clientNo + "] Parse error", parse);
                    sendFramed(out, gson.toJson(new JsonObjects.Error(parse)).getBytes(UTF_8));
                    continue;
                }

                try {
                    engine.layout(graph, new BasicProgressMonitor());
                    var sw = new java.io.StringWriter();
                    writeLayoutData(graph, sw, gson);
                    byte[] response = sw.toString().getBytes(UTF_8);
                    sendFramed(out, response);
                    LOG.info("[" + clientNo + "] Sent framed+gzip layout (" + response.length + " bytes uncompressed)");
                } catch (Exception layoutErr) {
                    LOG.log(Level.SEVERE, "[" + clientNo + "] Layout error", layoutErr);
                    sendFramed(out, gson.toJson(new JsonObjects.Error(layoutErr)).getBytes(UTF_8));
                }
            }
        } else { // LEGACY
            ElkNode graph;
            while ((graph = readGraph(in)) != null) {
                LOG.info("[" + clientNo + "] Received legacy JSON graph");
                try {
                    engine.layout(graph, new BasicProgressMonitor());
                    var writer = new java.io.OutputStreamWriter(out, UTF_8);
                    writeLayoutData(graph, writer, gson);
                    writer.flush();
                    LOG.info("[" + clientNo + "] Sent legacy JSON layout");
                } catch (Exception exc) {
                    var writer = new java.io.OutputStreamWriter(out, UTF_8);
                    gson.toJson(new JsonObjects.Error(exc), writer);
                    writer.flush();
                    LOG.log(Level.SEVERE, "[" + clientNo + "] Layout error", exc);
                }
            }
        }
    }

    private enum Mode { FRAMED, LEGACY }

    private Mode detectMode(BufferedInputStream in) throws Exception {
        in.mark(HEADER_LEN);
        byte[] probe = in.readNBytes(HEADER_LEN);
        if (probe.length == 0) {
            return Mode.LEGACY;
        }
        if (probe.length < HEADER_LEN) {
            in.reset();
            return Mode.LEGACY;
        }
        boolean magic = probe[0] == MAGIC[0] &&
                        probe[1] == MAGIC[1] &&
                        probe[2] == MAGIC[2] &&
                        probe[3] == MAGIC[3];
        in.reset();
        return magic ? Mode.FRAMED : Mode.LEGACY;
    }

    private byte[] readFramedRequest(BufferedInputStream in) throws Exception {
        // Attempt to read header; if EOF return null
        in.mark(HEADER_LEN);
        byte[] header = in.readNBytes(HEADER_LEN);
        if (header.length == 0) return null;
        if (header.length < HEADER_LEN) throw new EOFException("Truncated header");
        boolean magic = header[0] == MAGIC[0] && header[1] == MAGIC[1] && header[2] == MAGIC[2] && header[3] == MAGIC[3];
        if (!magic) {
            throw new IllegalStateException("Expected magic header for framed request");
        }
        int length = ((header[4] & 0xFF) << 24) |
                     ((header[5] & 0xFF) << 16) |
                     ((header[6] & 0xFF) << 8) |
                     (header[7] & 0xFF);
        if (length < 0) throw new IllegalArgumentException("Negative length");
        byte[] comp = in.readNBytes(length);
        if (comp.length != length) throw new EOFException("Truncated compressed payload");
        try (GZIPInputStream gis = new GZIPInputStream(new ByteArrayInputStream(comp))) {
            return gis.readAllBytes();
        }
    }

    private void sendFramed(OutputStream out, byte[] rawJson) throws Exception {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try (GZIPOutputStream gzos = new GZIPOutputStream(baos)) {
            gzos.write(rawJson);
        }
        byte[] comp = baos.toByteArray();
        byte[] header = new byte[HEADER_LEN];
        System.arraycopy(MAGIC, 0, header, 0, MAGIC.length);
        int len = comp.length;
        header[4] = (byte)(len >>> 24);
        header[5] = (byte)(len >>> 16);
        header[6] = (byte)(len >>> 8);
        header[7] = (byte)(len);
        out.write(header);
        out.write(comp);
        out.flush();
    }

    public static void main(String[] args) {
        new GzipFramedSocketElkServer().start();
    }
}